from typing import Dict, List, Optional, Set
from dataclasses import dataclass
from enum import Enum
from datetime import datetime

class TaskStatus(Enum):
    """Enumeration of possible task states."""
    PENDING = 'pending'
    RUNNING = 'running'
    COMPLETED = 'completed'
    FAILED = 'failed'

@dataclass
class TaskInfo:
    """Data class containing task information."""
    task_id: str
    worker_id: Optional[str]
    status: TaskStatus
    dependencies: Set[str]
    start_time: Optional[datetime]
    end_time: Optional[datetime]
    memory_usage: float  # in MB
    output_size: float   # in MB

class RuntimeInfo:
    """Manages runtime information for the distributed system.
    
    This class tracks task status, data dependencies, and resource usage
    across all workers in the system.
    """
    
    def __init__(self):
        self.tasks: Dict[str, TaskInfo] = {}
        self.worker_status: Dict[str, bool] = {}  # True if worker is active
        self.data_locations: Dict[str, List[str]] = {}  # data_id -> list of worker_ids
        self.memory_usage: Dict[str, float] = {}  # worker_id -> memory usage in MB
    
    def register_task(self, task_id: str, dependencies: Set[str]) -> None:
        """Register a new task in the system.
        
        Args:
            task_id: Unique identifier for the task
            dependencies: Set of task IDs this task depends on
        """
        self.tasks[task_id] = TaskInfo(
            task_id=task_id,
            worker_id=None,
            status=TaskStatus.PENDING,
            dependencies=dependencies,
            start_time=None,
            end_time=None,
            memory_usage=0.0,
            output_size=0.0
        )
    
    def update_task_status(self, task_id: str, status: TaskStatus,
                          worker_id: Optional[str] = None) -> None:
        """Update the status of a task.
        
        Args:
            task_id: Task identifier
            status: New status of the task
            worker_id: Worker assigned to the task (if any)
        """
        if task_id not in self.tasks:
            raise KeyError(f"Task {task_id} not found")
            
        task = self.tasks[task_id]
        task.status = status
        
        if status == TaskStatus.RUNNING:
            task.worker_id = worker_id
            task.start_time = datetime.now()
        elif status in [TaskStatus.COMPLETED, TaskStatus.FAILED]:
            task.end_time = datetime.now()
    
    def update_resource_usage(self, worker_id: str, memory_usage: float) -> None:
        """Update resource usage information for a worker.
        
        Args:
            worker_id: Worker identifier
            memory_usage: Current memory usage in MB
        """
        self.memory_usage[worker_id] = memory_usage
    
    def register_data_location(self, data_id: str, worker_id: str) -> None:
        """Register where data is stored in the system.
        
        Args:
            data_id: Identifier for the data
            worker_id: Worker where the data is stored
        """
        if data_id not in self.data_locations:
            self.data_locations[data_id] = []
        if worker_id not in self.data_locations[data_id]:
            self.data_locations[data_id].append(worker_id)
    
    def get_task_dependencies(self, task_id: str) -> Set[str]:
        """Get the dependencies for a task.
        
        Args:
            task_id: Task identifier
            
        Returns:
            Set[str]: Set of task IDs this task depends on
        """
        return self.tasks[task_id].dependencies if task_id in self.tasks else set()
    
    def get_data_locations(self, data_id: str) -> List[str]:
        """Get the locations where specific data is stored.
        
        Args:
            data_id: Data identifier
            
        Returns:
            List[str]: List of worker IDs where the data is stored
        """
        return self.data_locations.get(data_id, [])
    
    def get_worker_load(self, worker_id: str) -> float:
        """Get the current memory load of a worker.
        
        Args:
            worker_id: Worker identifier
            
        Returns:
            float: Current memory usage in MB
        """
        return self.memory_usage.get(worker_id, 0.0)
    
    def get_running_tasks(self) -> List[str]:
        """Get list of currently running tasks.
        
        Returns:
            List[str]: List of task IDs that are currently running
        """
        return [task_id for task_id, task in self.tasks.items()
                if task.status == TaskStatus.RUNNING]
    
    def get_completed_tasks(self) -> List[str]:
        """Get list of completed tasks.
        
        Returns:
            List[str]: List of task IDs that have completed
        """
        return [task_id for task_id, task in self.tasks.items()
                if task.status == TaskStatus.COMPLETED]
    
    def get_failed_tasks(self) -> List[str]:
        """Get list of failed tasks.
        
        Returns:
            List[str]: List of task IDs that have failed
        """
        return [task_id for task_id, task in self.tasks.items()
                if task.status == TaskStatus.FAILED]